{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "cell_0",
      "metadata": {},
      "source": [
        "# ContextGem: JsonObjectConcept Extraction with Examples"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cell_1",
      "metadata": {},
      "outputs": [],
      "source": [
        "%pip install -U contextgem"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "cell_2",
      "metadata": {},
      "source": [
        "To run the extraction, please provide your LLM details in the ``DocumentLLM(...)`` constructor further below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cell_3",
      "metadata": {},
      "outputs": [],
      "source": [
        "# ContextGem: JsonObjectConcept Extraction with Examples\n",
        "\n",
        "import os\n",
        "from pprint import pprint\n",
        "\n",
        "from contextgem import Document, DocumentLLM, JsonObjectConcept, JsonObjectExample\n",
        "\n",
        "\n",
        "# Document object with ambiguous medical report text\n",
        "medical_report = \"\"\"\n",
        "PATIENT ASSESSMENT\n",
        "Date: March 15, 2023\n",
        "Patient: John Doe (ID: 12345)\n",
        "\n",
        "Vital Signs:\n",
        "BP: 125/82 mmHg\n",
        "HR: 72 bpm\n",
        "Temp: 98.6\u00b0F\n",
        "SpO2: 98%\n",
        "\n",
        "Chief Complaint:\n",
        "Patient presents with persistent cough for 2 weeks, mild fever in evenings (up to 100.4\u00b0F), and fatigue. \n",
        "No shortness of breath. Patient reports recent travel to Southeast Asia 3 weeks ago.\n",
        "\n",
        "Assessment:\n",
        "Physical examination shows slight wheezing in upper right lung. No signs of pneumonia on chest X-ray.\n",
        "WBC slightly elevated at 11,500. Patient appears in stable condition but fatigued.\n",
        "\n",
        "Impression:\n",
        "1. Acute bronchitis, likely viral\n",
        "2. Rule out early TB given travel history\n",
        "3. Fatigue, likely secondary to infection\n",
        "\n",
        "Plan:\n",
        "- Rest for 5 days\n",
        "- Symptomatic treatment with over-the-counter cough suppressant\n",
        "- Follow-up in 1 week\n",
        "- TB test ordered\n",
        "\n",
        "Dr. Sarah Johnson, MD\n",
        "\"\"\"\n",
        "doc = Document(raw_text=medical_report)\n",
        "\n",
        "# Create a JsonObjectConcept for extracting medical assessment data\n",
        "# Without examples, the LLM might struggle with ambiguous fields or formatting variations\n",
        "medical_assessment_concept = JsonObjectConcept(\n",
        "    name=\"Medical Assessment\",\n",
        "    description=\"Key information from a patient medical assessment\",\n",
        "    structure={\n",
        "        \"patient\": {\n",
        "            \"id\": str,\n",
        "            \"vital_signs\": {\n",
        "                \"blood_pressure\": str,\n",
        "                \"heart_rate\": int,\n",
        "                \"temperature\": float,\n",
        "                \"oxygen_saturation\": int,\n",
        "            },\n",
        "        },\n",
        "        \"clinical\": {\n",
        "            \"symptoms\": list[str],\n",
        "            \"diagnosis\": list[str],\n",
        "            \"travel_history\": bool,\n",
        "        },\n",
        "        \"treatment\": {\"recommendations\": list[str], \"follow_up_days\": int},\n",
        "    },\n",
        "    # Examples provide helpful guidance on how to:\n",
        "    # 1. Map data from unstructured text to structured fields\n",
        "    # 2. Handle formatting variations (BP as \"120/80\" vs separate systolic/diastolic)\n",
        "    # 3. Extract implicit information (converting \"SpO2: 98%\" to just 98)\n",
        "    examples=[\n",
        "        JsonObjectExample(\n",
        "            content={\n",
        "                \"patient\": {\n",
        "                    \"id\": \"87654\",\n",
        "                    \"vital_signs\": {\n",
        "                        \"blood_pressure\": \"130/85\",\n",
        "                        \"heart_rate\": 68,\n",
        "                        \"temperature\": 98.2,\n",
        "                        \"oxygen_saturation\": 99,\n",
        "                    },\n",
        "                },\n",
        "                \"clinical\": {\n",
        "                    \"symptoms\": [\"headache\", \"dizziness\", \"nausea\"],\n",
        "                    \"diagnosis\": [\"Migraine\", \"Dehydration\"],\n",
        "                    \"travel_history\": False,\n",
        "                },\n",
        "                \"treatment\": {\n",
        "                    \"recommendations\": [\n",
        "                        \"Hydration\",\n",
        "                        \"Pain medication\",\n",
        "                        \"Dark room rest\",\n",
        "                    ],\n",
        "                    \"follow_up_days\": 14,\n",
        "                },\n",
        "            }\n",
        "        ),\n",
        "        JsonObjectExample(\n",
        "            content={\n",
        "                \"patient\": {\n",
        "                    \"id\": \"23456\",\n",
        "                    \"vital_signs\": {\n",
        "                        \"blood_pressure\": \"145/92\",\n",
        "                        \"heart_rate\": 88,\n",
        "                        \"temperature\": 100.8,\n",
        "                        \"oxygen_saturation\": 96,\n",
        "                    },\n",
        "                },\n",
        "                \"clinical\": {\n",
        "                    \"symptoms\": [\"sore throat\", \"cough\", \"fever\"],\n",
        "                    \"diagnosis\": [\"Strep throat\", \"Pharyngitis\"],\n",
        "                    \"travel_history\": True,\n",
        "                },\n",
        "                \"treatment\": {\n",
        "                    \"recommendations\": [\"Antibiotics\", \"Throat lozenges\", \"Rest\"],\n",
        "                    \"follow_up_days\": 7,\n",
        "                },\n",
        "            }\n",
        "        ),\n",
        "    ],\n",
        ")\n",
        "\n",
        "# Attach the concept to the document\n",
        "doc.add_concepts([medical_assessment_concept])\n",
        "\n",
        "# Configure DocumentLLM with your API parameters\n",
        "llm = DocumentLLM(\n",
        "    model=\"azure/gpt-4.1\",\n",
        "    api_key=os.getenv(\"CONTEXTGEM_AZURE_OPENAI_API_KEY\"),\n",
        "    api_version=os.getenv(\"CONTEXTGEM_AZURE_OPENAI_API_VERSION\"),\n",
        "    api_base=os.getenv(\"CONTEXTGEM_AZURE_OPENAI_API_BASE\"),\n",
        ")\n",
        "\n",
        "# Extract the concept from the document\n",
        "medical_assessment_concept = llm.extract_concepts_from_document(doc)[0]\n",
        "\n",
        "# Print the extracted medical assessment\n",
        "print(\"Extracted medical assessment:\")\n",
        "assessment = medical_assessment_concept.extracted_items[0].value\n",
        "pprint(assessment)\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}